package com.exercises;


import com.exercises.common.TreeNode;

import java.util.*;
import java.util.stream.Collectors;

class Exercise863 {
    Map<Integer, Set<Integer>> map = new HashMap<>();
    TreeNode target = null;
    int targetIndex = 0;

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        this.target = target;
        dfs(root, 0);
        Set<Integer> a = map.getOrDefault(targetIndex + k, new HashSet<>());
        Set<Integer> b = map.getOrDefault(Math.abs(targetIndex - k), new HashSet<>());
        a.addAll(b);
        return a.stream().filter(v -> v != target.val).collect(Collectors.toList());
    }

    public void dfs(TreeNode node, int index) {
        if (node == null) return;
        if (target.val == node.val) targetIndex = index;
        Set<Integer> list = map.getOrDefault(index, new HashSet<>());
        list.add(node.val);
        map.put(index, list);
        dfs(node.right, index + 1);
        dfs(node.left, index + 1);
    }
}

